package cn.smileyan.demos.core;

import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.state.ReducingState;
import org.apache.flink.api.common.state.ReducingStateDescriptor;
import org.apache.flink.api.common.typeutils.base.LongSerializer;
import org.apache.flink.streaming.api.windowing.triggers.Trigger;
import org.apache.flink.streaming.api.windowing.triggers.TriggerResult;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;

/**
 * 自定义 window 触发器
 * @author smileyan
 */
@Slf4j
public class CountAndTimeTrigger<T extends CountTrigger, W extends TimeWindow> extends Trigger<T, W> {
    /**
     * ReducingStateDescriptor 的 key 字段，上下文根据这个字段获取状态指
     */
    private static final String COUNT_KEY = "count";

    /**
     * ReducingStateDescriptor 根据聚合过程更新 count 结果
     */
    private final ReducingStateDescriptor<Long> stateDesc =
            new ReducingStateDescriptor<>(COUNT_KEY, new Sum(), LongSerializer.INSTANCE);

    @Override
    public TriggerResult onElement(T element, long timestamp, W window, TriggerContext ctx) throws Exception {
        ctx.registerEventTimeTimer(window.getEnd());
        ctx.registerProcessingTimeTimer(window.getEnd());
        final int size = element.triggerSize();
        final String id = element.triggerId();

        ReducingState<Long> count = ctx.getPartitionedState(stateDesc);
        count.add(1L);
        log.info("[{}] window: ({}, {}) -> merged {}", id, window.getStart(), window.getEnd(), count.get());

        if (count.get().intValue() == size) {
            log.info("[{} -> {}] merged successfully.", id, ctx.getCurrentWatermark());
            clear(window, ctx);
            ctx.getPartitionedState(stateDesc).clear();
            return TriggerResult.FIRE_AND_PURGE;
        } else if (count.get() > size) {
            log.warn("[{} -> {}] sent more than need {}", id, ctx.getCurrentWatermark(), size);
            return TriggerResult.PURGE;
        }

        return TriggerResult.CONTINUE;
    }

    @Override
    public TriggerResult onProcessingTime(long time, W window, TriggerContext ctx) throws Exception {
        if (time >= window.getEnd()) {
            log.debug("[ -> {}] onProcessingTime", ctx.getCurrentWatermark());
            return TriggerResult.FIRE_AND_PURGE;
        }
        return TriggerResult.CONTINUE;
    }

    @Override
    public TriggerResult onEventTime(long time, W window, TriggerContext ctx) throws Exception {
        if (time >= window.getEnd()) {
            log.debug("[ -> {}] onEventTime", ctx.getCurrentWatermark());
            return TriggerResult.FIRE_AND_PURGE;
        }
        return TriggerResult.CONTINUE;
    }

    @Override
    public void clear(W window, TriggerContext ctx) throws Exception {
        log.debug("[ -> {}] cleaning window ({}, {})", ctx.getCurrentWatermark(), window.getStart(), window.getEnd());
        ctx.deleteEventTimeTimer(window.getEnd());
        ctx.deleteProcessingTimeTimer(window.getEnd());
    }

    @Override
    public boolean canMerge() {
        return true;
    }

    @Override
    public void onMerge(TimeWindow window, OnMergeContext ctx) throws Exception {
        log.debug("[ -> {}] onMerge ({}, {})", ctx.getCurrentWatermark(), window.getStart(), window.getEnd());
        ctx.mergePartitionedState(stateDesc);
    }

    private static class Sum implements ReduceFunction<Long> {
        private static final long serialVersionUID = 1L;

        @Override
        public Long reduce(Long value1, Long value2) throws Exception {
            return value1 + value2;
        }
    }
}
